import math
import time
import numpy as np
from scipy.stats import norm
from multiprocessing import Process, Queue

# Genetic Algorithm
class GA:
  def __init__(self):
    self.nvar = 4
    self.varmin = [1e-3, 1e-3, -0.999, 1e-3]
    self.varmax = [1, 0.999, 0.999, 10]
    self.maxit = 2000
    self.npop = 100
    self.beta = 1
    self.pc = 1.0
    self.gamma = 0.1
    self.mu = 0.1
    self.sigma = 0.1
    self.epsilon = 1e-8
    self.update_ts = 0
    self.pop_position = []

  def update(self, training_data, queue):
    self.update_ts = time.time()
    nc = int(np.round(self.pc * self.npop / 2) * 2)
    T = []
    F = []
    strike_price = []
    market_iv = []
    for data in training_data:
      T.append(data[0])
      F.append(data[1])
      strike_price.append(data[2])
      market_iv.append(data[3])

    # Best Solution Ever Found
    bestsol = None
    bestsol_cost = np.inf

    # Initialize Population
    pop_cost = []
    first_time = True
    if len(self.pop_position) > 0:
       first_time = False
    for i in range(self.npop):
        if first_time:
          self.pop_position.append(np.random.uniform(self.varmin, self.varmax, self.nvar))
        pop_cost.append(self.fitness_func(self.pop_position[i], T, F, strike_price, market_iv))
        if pop_cost[i] < bestsol_cost:
            bestsol = self.pop_position[i]
            bestsol_cost = pop_cost[i]

    # Best Cost of Iterations
    bestcost = np.empty(self.maxit)

    # Main Loop
    for it in range(self.maxit):
        costs = np.array([x for x in pop_cost])
        avg_cost = np.mean(costs)
        if avg_cost != 0:
            costs = costs / avg_cost
        probs = np.exp(-self.beta * costs)
        popc = []
        popc_cost = []
        for _ in range(nc // 2):

            # Perform Roulette Wheel Selection
            p1 = self.pop_position[self.roulette_wheel_selection(probs)]
            p2 = self.pop_position[self.roulette_wheel_selection(probs)]

            # Perform Crossover
            c1, c2 = self.crossover(p1, p2, self.gamma)

            # Perform Mutation
            c1 = self.mutate(c1, self.mu, self.sigma)
            c2 = self.mutate(c2, self.mu, self.sigma)

            # Apply Bounds
            c1 = self.apply_bound(c1, self.varmin, self.varmax)
            c2 = self.apply_bound(c2, self.varmin, self.varmax)

            # Evaluate First Offspring
            c1_cost = self.fitness_func(c1, T, F, strike_price, market_iv)
            if c1_cost < bestsol_cost:
                bestsol_cost = c1_cost
                bestsol = c1

            # Evaluate Second Offspring
            c2_cost = self.fitness_func(c2, T, F, strike_price, market_iv)
            if c2_cost < bestsol_cost:
                bestsol_cost = c2_cost
                bestsol = c2

            # Add Offsprings to popc
            popc.append(c1)
            popc.append(c2)
            popc_cost.append(c1_cost)
            popc_cost.append(c2_cost)

        # Merge, Sort and Select
        self.pop_position += popc
        pop_cost += popc_cost
        idx = range(len(self.pop_position))
        idx = sorted(idx, key=lambda x: pop_cost[x]) 
        self.pop_position = [self.pop_position[i] for i in idx[0:self.npop]]
        pop_cost = [pop_cost[i] for i in idx[0:self.npop]]

        # Store Best Cost
        bestcost[it] = bestsol_cost
        if it >= 10 and bestcost[it - 10] - bestsol_cost <= self.epsilon:
          break
        
    print("mse:", bestsol_cost)
    print("update time:", time.time() - self.update_ts)
    queue.put([bestsol, self.pop_position])

  def crossover(self, p1, p2, gamma=0.1):
    alpha = np.random.uniform(-gamma, 1 + gamma, *p1.shape)
    c1 = alpha * p1 + (1 - alpha) * p2
    c2 = alpha * p2 + (1 - alpha) * p1
    return c1, c2


  def mutate(self, x, mu, sigma):
    y = x
    flag = np.random.rand(*x.shape) <= mu
    ind = np.argwhere(flag)
    y[ind] += sigma * np.random.randn(*ind.shape)
    return y


  def apply_bound(self, x, varmin, varmax):
    x = np.maximum(x, varmin)
    x = np.minimum(x, varmax)
    return x


  def roulette_wheel_selection(self, p):
    c = np.cumsum(p)
    r = sum(p) * np.random.rand()
    ind = np.argwhere(r <= c)
    return ind[0][0]
  
  # fit 90% points calculate mse
  def fitness_func(self, x, T, F, strike_prices, market_iv):
    alpha = x[0]
    beta = x[1]
    rho = x[2]
    nu = x[3]
    predicts = np.array([MathematicalTools.sabr_model(alpha, beta, rho, nu, T[i], F[i], strike_prices[i])
                         for i in range(len(strike_prices))])
    labels = np.array(market_iv)
    se = (predicts - labels)  ** 2
    length = int(len(se) * 0.9)
    mse = np.sum(np.sort(se)[:length]) / length
    return mse


class MathematicalTools:
  # Black Scholes model
  def bs_price(cp_flag, S, K, T, v):
    N = norm.cdf
    d1 = (math.log(S / K) + (v * v / 2.) * T) / (v * math.sqrt(T))
    d2 = d1 - v * math.sqrt(T)
    if cp_flag == 'C':
        price = S * N(d1) - K * N(d2)
    else:
        price = K * N(-d2) - S * N(-d1)
    return price

  # Derivative of price
  def bs_vega(S, K, T, v):
    n = norm.pdf
    d1 = (math.log(S / K) + (v * v / 2.) * T) / (v * math.sqrt(T))
    d2 = d1 - v * math.sqrt(T)
    return (S * n(d1) * (math.sqrt(T) / 2 - math.log(S / K) / (v * v * math.sqrt(T)))
            + K * n(d2) *(math.sqrt(T) / 2 + math.log(S / K) / (v * v * math.sqrt(T))))
  
  # Newton-Raphson method : calculate implied volatility 
  def find_vol(target_value, call_put, S, K, T):
    MAX_ITERATIONS = 100
    PRECISION = 1.0e-5
    sigma = 0.5
    for i in range(0, MAX_ITERATIONS):
        price = MathematicalTools.bs_price(call_put, S, K, T, sigma)
        vega = MathematicalTools.bs_vega(S, K, T, sigma)
        diff = target_value - price  # root
        if abs(diff) < PRECISION:
            return sigma
        if abs(vega) < PRECISION:
            return None
        sigma = sigma + diff / vega # f(x) / f'(x)
        if sigma > 1 - 1e-4 or sigma < 1e-4:
            return None
    return sigma
  
  # SABR model : calculate implied volatility without option price
  def sabr_model(alpha, beta, rho, nu, T, F, K, Model="Obloj2008"):
    if Model != "Obloj2008":
      z = nu / alpha * (F * K) ** ((1 - beta) / 2) * math.log(F / K)
      Xz = math.log((math.sqrt(1 - 2 * (rho * z) + z * z) + z - rho) / (1 - rho))
      sigma = (alpha * (1 + ((1 - beta) ** 2 / 24 * alpha ** 2 / (F * K) ** (1 - beta) +
                            alpha * beta * rho * nu / 4 / (F * K) ** ((1 - beta) / 2) +
                            (2 - 3 * rho ** 2) * nu ** 2 / 24) * T) /
              ((F * K) ** ((1 - beta) / 2) * (1 + (1 - beta) ** 2 / 24 * (math.log(F / K) ** 2) +
                                              (1 - beta) ** 4 / 1920 * math.log(F / K) ** 4)))
      if F != K:
          sigma *= z / Xz
    else:
      z = nu / alpha * (F ** (1 - beta) - K ** (1 - beta)) / (1 - beta)
      Xz = math.log((math.sqrt(1 - 2 * (rho * z) + z * z) + z - rho) / (1 - rho))
      sigma = nu * math.log(F / K) / Xz * (1 + ((1 - beta) ** 2 / 24 * alpha ** 2 / (F * K) ** (1 - beta) +
                            alpha * beta * rho * nu / 4 / (F * K) ** ((1 - beta) / 2) +
                            (2 - 3 * rho ** 2) * nu ** 2 / 24) * T)
    return sigma

class SharedValue:
  def __init__(self):
    self._forward_value = None
    self._bid_historical_data = dict()
    self._ask_historical_data = dict()
    self._bid_ga = GA()
    self._ask_ga = GA()
    self._bid_sample_num = 0
    self._ask_sample_num = 0
    self._bid_sample_num_update_ts = 0
    self._ask_sample_num_update_ts = 0
    self._bid_model_update_ts = 0
    self._ask_model_update_ts = 0
    self._bid_params = None
    self._ask_params = None
    self._bid_queue = Queue()
    self._ask_queue = Queue()
    self._bid_process = None
    self._ask_process = None


  def set_forward_value(self, forward_value):
    self._forward_value = forward_value
  
  def get_forward_value(self):
    return self._forward_value

  def update_bid_historical_data(self, key, value):
    self._bid_historical_data[key] = value
    training_data = []
    current_time = time.time()
    # use the historical data in 10 min
    for v in self._bid_historical_data.values(): 
      if v[0] > current_time and current_time < v[1] + 600:
        T = (v[0] - v[1]) / 31536000
        training_data.append([T, v[2], v[3], v[4]])
    sample_num = len(training_data)
    if sample_num != self._bid_sample_num:
       self._bid_sample_num = sample_num
       self._bid_sample_num_update_ts = current_time
       return
    if sample_num < 5:
       return
    if current_time > self._bid_sample_num_update_ts + 60 and (self._bid_process is None or not self._bid_process.is_alive()):  
      self._bid_model_update_ts = current_time
      self._bid_process = Process(target=self._bid_ga.update, args=(training_data, self._bid_queue,))
      self._bid_process.start()


  def update_ask_historical_data(self, key, value):
    self._ask_historical_data[key] = value
    training_data = []
    current_time = time.time()
    # use the historical data in 10 min
    for v in self._ask_historical_data.values():
      if v[0] > current_time and current_time < v[1] + 600:
        T = (v[0] - v[1]) / 31536000.
        training_data.append([T, v[2], v[3], v[4]])
    sample_num = len(training_data)
    if sample_num != self._ask_sample_num:
       self._ask_sample_num = sample_num
       self._ask_sample_num_update_ts = current_time
       return
    if sample_num < 5:
       return
    if current_time > self._ask_sample_num_update_ts + 60 and (self._ask_process is None or not self._ask_process.is_alive()):   
      self._ask_model_update_ts = current_time
      self._ask_process = Process(target=self._ask_ga.update, args=(training_data, self._ask_queue,))
      self._ask_process.start()
